import os
import pandas as pd
from pathlib import Path

# 1. 输入/输出文件夹
INPUT_DIR  = "Media_Monthly_Counts_Quad&Event"
OUTPUT_DIR = "Media_Monthly_Counts_Quad&Event (Added)"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 2. 构造完整月份索引（2014-01 至 2023-12，共 120 个月）&#8203;:contentReference[oaicite:2]{index=2}
FULL_MONTHS = pd.period_range("2014-01", "2023-12", freq="M")

def fill_and_fix(df: pd.DataFrame, is_event: bool=False) -> pd.DataFrame:
    """
    补齐 Month → FULL_MONTHS，并在 EventRootCode 表中
    若缺少 '20' 列则自动添加，所有值置 0。
    """
    df = df.copy()
    # 解析 Month 列为 PeriodIndex
    df['Month'] = pd.to_datetime(df['Month'], format="%Y-%m", errors='coerce') \
                    .dt.to_period('M')  # :contentReference[oaicite:3]{index=3}
    df = df.set_index('Month')

    # 数值化其余列，NaN→0
    df = df.apply(pd.to_numeric, errors='coerce').fillna(0)  # :contentReference[oaicite:4]{index=4}

    # 同月合并求和
    df = df.groupby(df.index).sum()  # :contentReference[oaicite:5]{index=5}

    # 补齐所有月份，fill 0
    df = df.reindex(FULL_MONTHS, fill_value=0)  # :contentReference[oaicite:6]{index=6}

    # 对 EventRootCode 表，若缺列 '20' 则添加
    if is_event and '20' not in df.columns:
        df['20'] = 0  # :contentReference[oaicite:7]{index=7}

    # 恢复 Month 列并格式化
    df = df.reset_index().rename(columns={'index':'Month'})
    df['Month'] = df['Month'].dt.strftime("%Y-%m")

    # 保证 Month 在最前面
    cols = ['Month'] + [c for c in df.columns if c!='Month']
    return df[cols]

# 3. 批量处理
for fname in os.listdir(INPUT_DIR):
    if not fname.lower().endswith(".xlsx"):
        continue
    media = Path(fname).stem
    in_path  = os.path.join(INPUT_DIR, fname)
    out_path = os.path.join(OUTPUT_DIR, fname)

    xls = pd.ExcelFile(in_path, engine="openpyxl")  # :contentReference[oaicite:8]{index=8}
    quad = pd.read_excel(xls, sheet_name="QuadClass",      engine="openpyxl")
    evt  = pd.read_excel(xls, sheet_name="EventRootCode", engine="openpyxl")

    quad_filled = fill_and_fix(quad, is_event=False)
    evt_filled  = fill_and_fix(evt,  is_event=True)

    with pd.ExcelWriter(out_path, engine="openpyxl") as writer:  # :contentReference[oaicite:9]{index=9}
        quad_filled.to_excel(writer, sheet_name="QuadClass",      index=False)  # :contentReference[oaicite:10]{index=10}
        evt_filled .to_excel(writer, sheet_name="EventRootCode",  index=False)  # :contentReference[oaicite:11]{index=11}

    print(f"✔ Processed: {media} → saved to {out_path}")
